"""
 python3.9 /data/x/app/moo/./moo/script/sg.py 60000 60001
"""
import sys
import threading
from threading import Event
from py4j.java_gateway import GatewayClient, JavaGateway, java_import, CallbackServerParameters

event = Event()

def process(jgw, params):
    dictParams = dict(params)

    modelName = "./moo/moo/huggingFaceTextClassification"
    text = dictParams.pop("text")
    print(modelName)
    print(text)
    cats = processRow(modelName, text, dictParams)
    print(len(cats))
    resultList = jgw.jvm.java.util.LinkedList()
    for i in range(len(cats)):
        row = jgw.jvm.java.util.HashMap()
        print(cats)
        row.put("label",cats[i]["label"])
        row.put("prob", cats[i]["score"])
        resultList.add(row)
    return resultList

from transformers import pipeline

def processRow(modelPath, sentence, params):
    device = -1
    if 'device' in params:
        device = params['device']
    classifier = pipeline("text-classification",model=modelPath, device=device)
    return classifier(sentence, top_k=int(params['top_k']))

class MooService(object):

    jgw = None

    def __init__(self, jgw):
          self.jgw = jgw

    def call(self, params):
        return process(self.jgw, params)

    def stop(self):
        event.set()

    def waitStop(self):
        event.wait()

    ##实现java的
    class Java:
        implements = ['org.datayoo.sengee.service.moo.MooService']

def main(argv):
    client = GatewayClient(address="127.0.0.1", port=int(argv[1]))
    csp = CallbackServerParameters(port=int(argv[2])) #回调python
    jgw = JavaGateway(client, start_callback_server=True, callback_server_parameters=csp)
    java_import(jgw.jvm,"java.util.*")


    pisvc = jgw.entry_point
    service = MooService(jgw)
    # entry_point对应java的MooInteractor，所以可执行registMooService方法，这样java端调用，就是动态代理调用了python
    pisvc.registMooService(service)
    print('%s is registed!'%  argv[1]) # moo节点的端口
    service.waitStop()
    print('%s is stopped!' % argv[1]) #moo节点节点关闭，这里应该就停止了；
    jgw.close(close_callback_server_connections=True)



if __name__ == "__main__":
   main(sys.argv)